import numpy as np
import cv2
import torchvision.transforms as transforms
from PIL import Image  # 用于处理图片
import torch
import joblib  # 用于加载 SVM 模型
import mss
import time
import threading
from queue import Queue

# 常量
STD_LIST = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, "+", '-', ".", '*', "(", ")", '/']
SCALE_FACTOR = 0.5  # 缩放因子

# 加载 SVM 模型
svm_model = joblib.load('mnist_svm.pkl')

# 图像转换
transform = transforms.Compose([
    transforms.Resize((28, 28)),  # 改为适合SVM输入
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # MNIST normalization values
])


def accessBinary(img, threshold=128):
    img = 255 - img  # 反转颜色
    kernel = np.ones((3, 3), np.uint8)
    img = cv2.dilate(img, kernel, iterations=1)
    _, img = cv2.threshold(img, threshold, 0, cv2.THRESH_TOZERO)
    return img


def showResults(img, borders, results=None):
    if borders is None:
        return
    # 将结果按照边框的横坐标排序
    res = ""  # 最后的式子
    sorted_results = sorted(zip(borders, results), key=lambda x: x[0][0][0])
    for i, (border, result) in enumerate(sorted_results):
        cv2.rectangle(img, border[0], border[1], (0, 0, 255))
        if result is not None:
            cv2.putText(img, str(result), border[0], cv2.FONT_HERSHEY_COMPLEX, 0.8, (0, 255, 0), 1)
            res += str(STD_LIST[result])
    try:
        print(res)
        res = res[:len(res) - 3]
        print(res + " = " + str(eval(res)))
    except:
        print("FAKE: " + res)

    img_resized = cv2.resize(img, (int(img.shape[1] * SCALE_FACTOR), int(img.shape[0] * SCALE_FACTOR)))
    cv2.imshow('Real-time Digit Recognition', img_resized)
    cv2.waitKey(1)


def transMNIST(img, borders, size=(28, 28)):
    imgData = np.zeros((len(borders), size[0], size[0], 1), dtype='uint8')
    img = accessBinary(img)
    for i, border in enumerate(borders):
        borderImg = img[border[0][1]:border[1][1], border[0][0]:border[1][0]]
        extendPixel = (max(borderImg.shape) - min(borderImg.shape)) // 2
        targetImg = cv2.copyMakeBorder(borderImg, 7, 7, extendPixel + 7, extendPixel + 7, cv2.BORDER_CONSTANT)
        targetImg = cv2.resize(targetImg, size)
        targetImg = np.expand_dims(targetImg, axis=-1)
        imgData[i] = targetImg
    return imgData


def findBorderContours(img, maxArea=50):
    img = accessBinary(img)
    contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    borders = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        if w * h > maxArea:
            border = [(x, y), (x + w, y + h)]
            borders.append(border)
    return borders


def preprocess_image(image):
    image = Image.fromarray(image.squeeze())  # NumPy数组转换为PIL图像
    image = transform(image)
    image = image.view(-1)  # 展平成为一维
    return image.numpy()


def predict(model, image):
    input_data = preprocess_image(image)
    pred_label = model.predict([input_data])[0]  # 预测结果
    return pred_label


def capture_screen(queue, monitor):
    sct = mss.mss()
    while True:
        screenshot = np.array(sct.grab(monitor))
        queue.put(screenshot)
        time.sleep(0.1)  # 控制帧率


def process_frame(model, frame_queue):
    last_recognition_time = time.time()
    recognition_interval = 1  # 2-3 seconds interval

    while True:
        if not frame_queue.empty():
            screenshot = frame_queue.get()
            screenshot = cv2.cvtColor(screenshot, cv2.COLOR_BGRA2BGR)
            current_time = time.time()

            if current_time - last_recognition_time >= recognition_interval:
                gray_screenshot = cv2.cvtColor(screenshot, cv2.COLOR_BGR2GRAY)
                borders = findBorderContours(gray_screenshot)
                imgData = transMNIST(gray_screenshot, borders)
                recognize_results = []
                for i in range(len(borders)):
                    result = predict(model, imgData[i])
                    recognize_results.append(result)

                # 按照横坐标排序边框和结果
                sorted_indexes = sorted(range(len(borders)), key=lambda i: borders[i][0][0])
                sorted_borders = [borders[i] for i in sorted_indexes]
                sorted_results = [recognize_results[i] for i in sorted_indexes]

                showResults(screenshot, sorted_borders, sorted_results)

                last_recognition_time = current_time
            else:
                showResults(screenshot, borders=None)


def real_time_detection():
    frame_queue = Queue();
    monitor = {"top": 54, "left": 8, "width": 1895, "height": 962}  # 修改为需要的显示器区域

    screen_thread = threading.Thread(target=capture_screen, args=(frame_queue, monitor))
    process_thread = threading.Thread(target=process_frame, args=(svm_model, frame_queue))

    screen_thread.start()
    process_thread.start()

    screen_thread.join()
    process_thread.join()

real_time_detection()
